package com.yufeng.collaboration.handler;

import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class MyWebSocketHandler extends TextWebSocketHandler {

    // 用户 ID 与 WebSocketSession 的映射
    private final Map<Long, WebSocketSession> userSessionMap = new HashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 假设前端在建立连接时传递了用户 ID 参数
        Long userId = (Long) session.getAttributes().get("userId");
        if (userId != null) {
            userSessionMap.put(userId, session);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        // 处理普通文本消息的逻辑，可根据需求扩展
        for (WebSocketSession webSocketSession : userSessionMap.values()) {
            if (webSocketSession.isOpen() && !webSocketSession.getId().equals(session.getId())) {
                webSocketSession.sendMessage(message);
            }
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        // 移除断开连接的用户的映射
        userSessionMap.values().remove(session);
    }

    /**
     * 向指定用户发送消息
     * @param userId 用户 ID
     * @param message 消息内容
     */
    public void sendMessageToUser(Long userId, String message) {
        WebSocketSession session = userSessionMap.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(message));
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}    